# 引入必要的库
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, auc, precision_recall_curve, average_precision_score
from keras.utils import to_categorical
from scipy import interp

target_name_dict = {'zhengchang': 'norm', 'zhadong': 'defect_1',
                    'maoban': 'defect_2', 'cadong': 'defect_3',
                    'maodong': 'defect_4', 'zhixi': 'defect_5',
                    'diaojing': 'defect_6', 'quejing': 'defect_7',
                    'tiaohua': 'defect_8', 'youzi': 'defect_9',
                    'wuzi': 'defect_9', 'qita': 'defect_10',
                    'xiaci': 'defect'}

name_target_dict = dict(zip(target_name_dict.values(), target_name_dict.keys()))


def cal_auc(y_pred, y_true, columns, draw=False):
    """
    计算roc
    :param y_pred: 预测值 onehot编码
    :param y_true: 真实值 onehot编码
    :param columns: 每一列所代表的类别名称
    :return: 各roc值,为一字典，keys='macro','micro',0~len(y_pred[0])
    """
    # 计算每一类的ROC
    fpr = dict()
    tpr = dict()
    roc_auc = dict()
    for i in range(len(y_pred[0])):
        fpr[i], tpr[i], _ = roc_curve(y_true[:, i], y_pred[:, i])
        key = name_target_dict[columns[i]]
        roc_auc[key] = round(auc(fpr[i], tpr[i]), 4)

    # micro-average ROC(方法二，转化为二分类法）
    # fpr["micro"], tpr["micro"], _ = roc_curve(y_true.ravel(), y_pred.ravel())
    # roc_auc["micro"] = round(auc(fpr["micro"], tpr["micro"]), 4)

    # macro-average ROC(方法一,平均法）
    all_fpr = np.unique(np.concatenate([fpr[i] for i in range(len(y_pred[0]))]))
    mean_tpr = np.zeros_like(all_fpr)
    for i in range(len(y_pred[0])):
        mean_tpr += interp(all_fpr, fpr[i], tpr[i])
    mean_tpr /= len(y_pred[0])
    fpr["macro"] = all_fpr
    tpr["macro"] = mean_tpr
    roc_auc["macro"] = round(auc(fpr["macro"], tpr["macro"]), 4)

    if draw:
        # 画两类ROC曲线
        plt.figure(figsize=(10, 10))
        # plt.plot(fpr["micro"], tpr["micro"],
        #          label='micro-average ROC curve (area = {0:0.2f})'.format(roc_auc["micro"]), linestyle=':', linewidth=4)

        plt.plot(fpr["macro"], tpr["macro"],
                 label='macro-average ROC curve (area = {0:0.2f})'.format(roc_auc["macro"]), linestyle=':', linewidth=4)

        # 画所有分类别的ROC曲线
        for i in range(len(y_pred[0])):
            key = name_target_dict[columns[i]]
            plt.plot(fpr[i], tpr[i], lw=2, label='ROC of {0} (area = {1:0.2f})'.format(key, roc_auc[key]))

        plt.plot([0, 1], [0, 1], 'k--', lw=1, alpha=0.5)
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('roc_curve')
        plt.legend(loc="lower right")
        plt.show()
    return roc_auc


def cal_mAP(y_pred, y_true, columns, draw=False):
    """
    计算roc
    :param y_pred: 预测值 onehot编码
    :param y_true: 真实值 onehot编码
    :param columns: 每一列所代表的类别名称
    :return: 各ap值,为一字典，keys='macro','micro',0~len(y_pred[0])
    """
    # 计算每一类的AP
    precision = dict()
    recall = dict()
    ap = dict()
    for i in range(len(y_pred[0])):
        precision[i], recall[i], _ = precision_recall_curve(y_true[:, i], y_pred[:, i])
    for index, value in enumerate(average_precision_score(y_true, y_pred, average=None)):
        key = name_target_dict[columns[index]]
        ap[key] = round(value, 4)

    # micro-average mAP(方法二，转化为二分类法）
    # precision["micro"], recall["micro"], _ = precision_recall_curve(y_true.ravel(), y_pred.ravel())
    # ap['micro'] = round(average_precision_score(y_true, y_pred, average='micro'), 4)

    # macro-average ROC(方法一,平均法）
    all_recall = np.unique(np.concatenate([recall[i] for i in range(len(y_pred[0]))]))
    mean_precision = np.zeros_like(all_recall)
    for i in range(len(y_pred[0])):
        mean_precision += interp(all_recall, recall[i], precision[i])
    mean_precision /= len(y_pred[0])
    recall["macro"] = all_recall
    precision["macro"] = mean_precision
    ap['macro'] = round(average_precision_score(y_true, y_pred, average='macro'), 4)

    if draw:
        # 画两类mAP曲线
        plt.figure(figsize=(10, 10))
        # plt.plot(recall["micro"], precision["micro"],
        #          label='micro-average mAP curve (area = {0:0.2f})'.format(ap["micro"]), linestyle=':', linewidth=4)

        # plt.plot(recall["macro"], precision["macro"],
        #          label='macro-average mAP curve (area = {0:0.2f})'.format(ap["macro"]), linestyle=':', linewidth=4)

        # 画所有分类别的AP曲线
        for i in range(len(y_pred[0])):
            key = name_target_dict[columns[i]]
            plt.plot(recall[i], precision[i], lw=2, label='AP of {0} (area = {1:0.2f})'.format(key, ap[key]))

        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('Recall')
        plt.ylabel('Precision')
        plt.title('precision_recall_curve')
        plt.legend(loc="upper right")
        plt.show()

    return ap


def cal_acc(y_pred, y_true):
    cnt = 0
    for i in range(len(y_pred)):
        if y_pred[i].argmax() == y_true[i].argmax():
            cnt += 1
    return round(cnt / len(y_pred), 4)


def main(csv_path, draw=False):
    data = pd.read_csv(csv_path)
    y_pred = data.iloc[:, 1:-1]
    y_pred = np.array(y_pred)

    y_true = data.loc[:, 'real']
    columns = list(data.iloc[:, 1:-1].columns)
    y_true = [columns.index(label.strip()) for label in y_true]
    y_true = to_categorical(y_true, num_classes=len(columns))
    y_true = np.array(y_true)

    auc = cal_auc(y_pred, y_true, columns, draw=draw)
    mAP = cal_mAP(y_pred, y_true, columns, draw=draw)
    acc = cal_acc(y_pred, y_true)

    print('auc', auc)
    print('mAP', mAP)
    # print('score[micro]', "{:.4f}".format(0.7 * auc['micro'] + 0.3 * mAP['micro']))
    print('score[macro]', "{:.4f}".format(0.7 * auc['macro'] + 0.3 * mAP['macro']))
    print('acc', "{:.4f}".format(acc))


if __name__ == '__main__':
    main('../submit/balance_equalize_merge_bin.csv', draw=True)
    # main('../submit/submit_model_weight.10-0.7704_balance.hdf5.csv')
    # main('../submit/submit_model_weight.12-0.8316.hdf5.csv')
    # main('../submit/submit_model_weight.bin.09-0.8827.hdf5.csv')
